#pragma OPENCL EXTENSION cl_amd_printf : enable
#define N_TYPE          1     //by mako
#define MAX_NEIGH       128   //by mako
#define N_UNROLL        4     //unroll, future by mako
#define SMALL           1.E-14f //avoid dividing by 0.

//add_f_e_slow: >O(N^2),all,simple for test 
//add_f_e_all: O(N),all
//add_f_e_sgroup: O(N),static-group
//add_f_e_dgroup: O(N),dynamic-group (mako)

//virial
//E= 4*epsi*((sig/r)^12 - (sig/r)^6). Note: the prefactor 4!
__kernel void add_f_e_vir_all( 
    const int n_atoms,
    const float4 box_l,  //boxsize, s3:should be no-zero to avoid dividing by 0
    const float4 box_p,  //box periodic, 1:periodic, 0:nonperiodic
    __global float4 *r_t,  //[n_atoms] position, type
    __global float4 *f_e,  //[n_atoms]force, en
    __global float8 *vir,  //[n_atoms]virial
    const __global float4 *par,  //[N_TYPE][N_TYPE] parameters s0:epsi,s1:sig^2,s2:rcut^2
    __global int *neigh,   //[MAX_NEIGH][n_atoms] neighbor list, <0 means NULL
    __global int *n_neigh  //[n_atoms] real number of neighbor
    ) 
{
    int iglb = get_global_id(0); 
    int sglb = get_global_size(0);

    int i = iglb;
    while(i<n_atoms){
        //cal f_i
        int ni = n_neigh[i];
        float4 xi = r_t[i];
        float4 fi0 = f_e[i];
        //--prefetch
        int j_next = neigh[i];  //j_next = (j_next<0)? i : j_next;
        float4 xj_next = r_t[j_next];
        //--end prefetch
        int pidx0 = xi.w*N_TYPE;
        float4 fi = 0.f;
        int pidx;
        float4 xj; 
        float epsi,sig2,r2cut;
        float dx,dy,dz,r2,r2inv,sor2,sor6,fr;
        int jj = 0;
        while(jj<ni) {//for(int kk=0;kk<N_UNROLL;kk++){ //let compiler do unrolling 
        //manual unroll:
            xj = xj_next;
            //xj = r_t[j_next];
            //--prefetch
            j_next = neigh[i+(jj+1)*n_atoms]; 
            j_next = (j_next<0) ? i : j_next;
            xj_next = r_t[j_next];
            //--end prefetch
            pidx = pidx0 + xj.w; //parametr index
            epsi  = par[pidx].s0; //epsilon
            sig2  = par[pidx].s1; //sigma^2
            r2cut = par[pidx].s2; //rcut^2
            //
            dx = xj.x - xi.x;
            dy = xj.y - xi.y;
            dz = xj.z - xi.z;
            dx -= round(dx/box_l.x)*box_l.x*box_p.x; //periodic or nonperiodic box
            dy -= round(dy/box_l.y)*box_l.y*box_p.y; //periodic or nonperiodic box
            dz -= round(dz/box_l.z)*box_l.z*box_p.z; //periodic or nonperiodic box
            r2 = dx*dx + dy*dy + dz*dz;
            //NOTE: the judgement of r2>r2cut is needed if using rcut+skin
            r2inv = (r2>r2cut || r2<SMALL)? 0.f : 1.f/r2; //avoid dividing 0
            //float r2inv = (r2>r2cut)? 0.f : 1.f/(r2+SMALL); //avoid dividing 0
            sor2 = sig2*r2inv;
            sor6 = sor2*sor2*sor2;
            fr = epsi*sor6 * (24.f - 48.f*sor6);
            //
            fi.x += fr*r2inv * dx;
            fi.y += fr*r2inv * dy;
            fi.z += fr*r2inv * dz;
            fi.w += epsi*2.f*sor6*(sor6-1.f); //en_i
            ++jj;

            xj = xj_next;
            //xj = r_t[j_next];
            //--prefetch
            j_next = neigh[i+(jj+1)*n_atoms]; 
            j_next = (j_next<0) ? i : j_next;
            xj_next = r_t[j_next];
            //--end prefetch
            pidx = pidx0 + xj.w; //parametr index
            epsi  = par[pidx].s0; //epsilon
            sig2  = par[pidx].s1; //sigma^2
            r2cut = par[pidx].s2; //rcut^2
            //
            dx = xj.x - xi.x;
            dy = xj.y - xi.y;
            dz = xj.z - xi.z;
            dx -= round(dx/box_l.x)*box_l.x*box_p.x; //periodic or nonperiodic box
            dy -= round(dy/box_l.y)*box_l.y*box_p.y; //periodic or nonperiodic box
            dz -= round(dz/box_l.z)*box_l.z*box_p.z; //periodic or nonperiodic box
            r2 = dx*dx + dy*dy + dz*dz;
            //NOTE: the judgement of r2>r2cut is needed if using rcut+skin
            r2inv = (r2>r2cut || r2<SMALL)? 0.f : 1.f/r2; //avoid dividing 0
            //float r2inv = (r2>r2cut)? 0.f : 1.f/(r2+SMALL); //avoid dividing 0
            sor2 = sig2*r2inv;
            sor6 = sor2*sor2*sor2;
            fr = epsi*sor6 * (24.f - 48.f*sor6);
            //
            fi.x += fr*r2inv * dx;
            fi.y += fr*r2inv * dy;
            fi.z += fr*r2inv * dz;
            fi.w += epsi*2.f*sor6*(sor6-1.f); //en_i
            ++jj;

            xj = xj_next;
            //xj = r_t[j_next];
            //--prefetch
            j_next = neigh[i+(jj+1)*n_atoms]; 
            j_next = (j_next<0) ? i : j_next;
            xj_next = r_t[j_next];
            //--end prefetch
            pidx = pidx0 + xj.w; //parametr index
            epsi  = par[pidx].s0; //epsilon
            sig2  = par[pidx].s1; //sigma^2
            r2cut = par[pidx].s2; //rcut^2
            //
            dx = xj.x - xi.x;
            dy = xj.y - xi.y;
            dz = xj.z - xi.z;
            dx -= round(dx/box_l.x)*box_l.x*box_p.x; //periodic or nonperiodic box
            dy -= round(dy/box_l.y)*box_l.y*box_p.y; //periodic or nonperiodic box
            dz -= round(dz/box_l.z)*box_l.z*box_p.z; //periodic or nonperiodic box
            r2 = dx*dx + dy*dy + dz*dz;
            //NOTE: the judgement of r2>r2cut is needed if using rcut+skin
            r2inv = (r2>r2cut || r2<SMALL)? 0.f : 1.f/r2; //avoid dividing 0
            //float r2inv = (r2>r2cut)? 0.f : 1.f/(r2+SMALL); //avoid dividing 0
            sor2 = sig2*r2inv;
            sor6 = sor2*sor2*sor2;
            fr = epsi*sor6 * (24.f - 48.f*sor6);
            //
            fi.x += fr*r2inv * dx;
            fi.y += fr*r2inv * dy;
            fi.z += fr*r2inv * dz;
            fi.w += epsi*2.f*sor6*(sor6-1.f); //en_i
            ++jj;

            xj = xj_next;
            //xj = r_t[j_next];
            //--prefetch
            j_next = neigh[i+(jj+1)*n_atoms]; 
            j_next = (j_next<0) ? i : j_next;
            xj_next = r_t[j_next];
            //--end prefetch
            pidx = pidx0 + xj.w; //parametr index
            epsi  = par[pidx].s0; //epsilon
            sig2  = par[pidx].s1; //sigma^2
            r2cut = par[pidx].s2; //rcut^2
            //
            dx = xj.x - xi.x;
            dy = xj.y - xi.y;
            dz = xj.z - xi.z;
            dx -= round(dx/box_l.x)*box_l.x*box_p.x; //periodic or nonperiodic box
            dy -= round(dy/box_l.y)*box_l.y*box_p.y; //periodic or nonperiodic box
            dz -= round(dz/box_l.z)*box_l.z*box_p.z; //periodic or nonperiodic box
            r2 = dx*dx + dy*dy + dz*dz;
            //NOTE: the judgement of r2>r2cut is needed if using rcut+skin
            r2inv = (r2>r2cut || r2<SMALL)? 0.f : 1.f/r2; //avoid dividing 0
            //float r2inv = (r2>r2cut)? 0.f : 1.f/(r2+SMALL); //avoid dividing 0
            sor2 = sig2*r2inv;
            sor6 = sor2*sor2*sor2;
            fr = epsi*sor6 * (24.f - 48.f*sor6);
            //
            fi.x += fr*r2inv * dx;
            fi.y += fr*r2inv * dy;
            fi.z += fr*r2inv * dz;
            fi.w += epsi*2.f*sor6*(sor6-1.f); //en_i
            ++jj;
        } //end-while
        f_e[i] = fi0 + fi;
        //f_e[i] = fi;
        //if(ig<4) printf("ig:%d,fi:%v4f,f_e:%v4f\n",ig,fi,f_e[i]);
        ///
        i += sglb;
    }
}
//----------------------------------------------------

//----------------------------------------------------
//E= 4*epsi*((sig/r)^12 - (sig/r)^6). Note: the prefactor 4!
__kernel void add_f_e_slow( 
    const int n_atoms,
    const float4 box_l,  //boxsize, s3:should be no-zero to avoid dividing by 0
    const float4 box_p,  //box periodic
    __global float4 *r_t,  //position, type
    __global float4 *f_e,  //force, en
    __global float8 *vir,  //virial
    __global float4 *par   //parameters s0:epsi, s1:sig, s2:r2_cut
    ) 
{
    int iglb = get_global_id(0); 
    int sglb = get_global_size(0);

    int i = iglb;
    while(i<n_atoms){
        //cal f_i
        float4 xi = r_t[i];
        float4 fi = f_e[i]; //prefetch
        float4 xj,p,d;
        int idx;
        float r2,r2inv,sor2,r6inv,fr;
        int j = 0;
        //while(j<400){ //for test
        while(j<n_atoms){//for(int kk=0;kk<N_UNROLL;kk++){//hope compiler unroll 
        //manual unroll
            xj = r_t[j];
            idx = xi.w*N_TYPE + xj.w; //type index
            p = par[idx]; 
            d = xj - xi;
            d = d - round(d/box_l)*box_l*box_p; //periodic or nonperiodic box
            r2 = d.x*d.x + d.y*d.y + d.z*d.z;
            r2inv = (r2<SMALL || r2>p.s2 )? 0.f : 1.f/r2; //avoiding self.
            sor2 = p.s1*r2inv;
            r6inv = sor2 * sor2 * sor2;
            fr = p.s0*r6inv * (24.f - 48.f*r6inv);
            fi.x += fr * r2inv * d.x;
            fi.y += fr * r2inv * d.y;
            fi.z += fr * r2inv * d.z;
            fi.w += p.s0*2.f*r6inv*(r6inv-1.f); //en_i
            j++;
            //
        } //end-while
        f_e[i] += fi;
        //if(ig<4) printf("ig:%d,fi:%v4f,f_e:%v4f\n",ig,fi,f_e[i]);
        ///
        i += sglb;
    }
}
